"""
Unified event type definitions with progressive disclosure.

This module consolidates basic and advanced event functionality into a cohesive
system that supports both simple and complex use cases while following SOLID principles.
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import (
    Any,
    Dict,
    List,
    Optional,
    Union,
    TypeVar,
    Protocol,
    runtime_checkable,
    ClassVar,
    Self,
)
import time
import uuid
from datetime import datetime
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import pytorch_lightning as pl


# ========== Core Enumerations ==========


class EventType(Enum):
    """
    Unified hierarchical event type enumeration.

    Uses dot notation for better organization and scalability.
    Maintains compatibility mappings for simple naming.
    """

    # === System Lifecycle Events ===
    SYSTEM_INITIALIZED = "system.initialized"
    SYSTEM_READY = "system.ready"
    SYSTEM_SHUTDOWN = "system.shutdown"

    # === Task Management Events ===
    TASK_STARTED = "task.started"  # was: task_start
    TASK_COMPLETED = "task.completed"  # was: task_end
    TASK_SWITCHED = "task.switched"  # was: task_switch
    TASK_PREPARED = "task.prepared"  # was: task_prepared
    TASK_CONFIGURATION_CHANGED = "task.configuration_changed"

    # === Training Lifecycle Events ===
    TRAINING_STARTED = "training.started"  # was: training_start
    TRAINING_COMPLETED = "training.completed"  # was: training_end
    TRAINING_PAUSED = "training.paused"
    TRAINING_RESUMED = "training.resumed"

    EPOCH_STARTED = "epoch.started"  # was: epoch_start
    EPOCH_COMPLETED = "epoch.completed"  # was: epoch_end

    BATCH_STARTED = "batch.started"  # was: batch_start
    BATCH_COMPLETED = "batch.completed"  # was: batch_end
    BATCH_PREPROCESSED = "batch.preprocessed"

    STEP_STARTED = "step.started"  # was: step_start
    STEP_COMPLETED = "step.completed"  # was: step_end

    # === Strategy Request/Response Events ===
    LOSS_COMPUTATION_REQUESTED = (
        "strategy.loss_computation.requested"  # was: loss_compute_request
    )
    LOSS_COMPUTATION_COMPLETED = (
        "strategy.loss_computation.completed"  # was: loss_compute_response
    )

    BATCH_PREPROCESSING_REQUESTED = (
        "strategy.batch_preprocessing.requested"  # was: batch_preprocess_request
    )
    BATCH_PREPROCESSING_COMPLETED = (
        "strategy.batch_preprocessing.completed"  # was: batch_preprocess_response
    )

    MEMORY_UPDATE_REQUESTED = (
        "strategy.memory_update.requested"  # was: memory_update_request
    )
    MEMORY_UPDATE_COMPLETED = (
        "strategy.memory_update.completed"  # was: memory_update_response
    )

    PARAMETER_IMPORTANCE_UPDATE_REQUESTED = (
        "strategy.parameter_importance.requested"  # was: parameter_importance_update
    )
    PARAMETER_IMPORTANCE_UPDATE_COMPLETED = "strategy.parameter_importance.completed"

    STRATEGY_STATE_CHANGED = "strategy.state_changed"  # was: strategy_state_change
    STRATEGY_CONFIGURED = "strategy.configured"

    # === Model Events ===
    MODEL_FORWARD_PASS_STARTED = "model.forward_pass.started"
    MODEL_FORWARD_PASS_COMPLETED = "model.forward_pass.completed"
    MODEL_BACKWARD_PASS_STARTED = "model.backward_pass.started"
    MODEL_BACKWARD_PASS_COMPLETED = "model.backward_pass.completed"

    MODEL_WEIGHTS_UPDATED = "model.weights.updated"  # was: parameter_update
    MODEL_WEIGHTS_FROZEN = "model.weights.frozen"  # was: weights_frozen
    MODEL_WEIGHTS_UNFROZEN = "model.weights.unfrozen"  # was: weights_unfrozen

    MODEL_SAVED = "model.saved"  # was: model_saved
    MODEL_LOADED = "model.loaded"  # was: model_loaded
    MODEL_CHECKPOINTED = "model.checkpointed"

    GRADIENT_COMPUTED = "model.gradient.computed"  # was: gradient_computed

    # === Data Events ===
    DATASET_LOADED = "data.dataset.loaded"  # was: dataset_available
    DATASET_AUGMENTED = "data.dataset.augmented"
    DATALOADER_CREATED = "data.dataloader.created"

    BATCH_AVAILABLE = "data.batch.available"  # was: batch_available
    BATCH_FILTERED = "data.batch.filtered"  # was: data_filtered
    BATCH_AUGMENTED = "data.batch.augmented"  # was: data_augmented

    # === Evaluation Events ===
    EVALUATION_STARTED = "evaluation.started"  # was: evaluation_start
    EVALUATION_COMPLETED = "evaluation.completed"  # was: evaluation_end
    VALIDATION_STARTED = "validation.started"  # was: validation_start
    VALIDATION_COMPLETED = "validation.completed"  # was: validation_end
    TESTING_STARTED = "testing.started"
    TESTING_COMPLETED = "testing.completed"

    METRICS_COMPUTED = "evaluation.metrics.computed"  # was: metrics_computed
    METRICS_LOGGED = "evaluation.metrics.logged"

    # === Memory Management Events ===
    MEMORY_ALLOCATED = "memory.allocated"  # was: memory_allocated
    MEMORY_DEALLOCATED = "memory.deallocated"  # was: memory_freed
    MEMORY_PRESSURE_DETECTED = "memory.pressure.detected"  # was: memory_pressure
    MEMORY_CLEANUP_TRIGGERED = "memory.cleanup.triggered"

    # === Error and Warning Events ===
    ERROR_OCCURRED = "error.occurred"  # was: error_occurred
    WARNING_ISSUED = "warning.issued"  # was: warning_issued
    EXCEPTION_CAUGHT = "exception.caught"
    RECOVERY_ATTEMPTED = "recovery.attempted"

    # === Monitoring and Logging Events ===
    METRICS_UPDATED = "monitoring.metrics.updated"  # was: log_metrics
    PERFORMANCE_MEASURED = "monitoring.performance.measured"
    RESOURCE_USAGE_MEASURED = "monitoring.resource_usage.measured"

    LOG_MESSAGE_CREATED = "logging.message.created"
    CHECKPOINT_SAVED = "logging.checkpoint.saved"  # was: checkpoint_saved
    CHECKPOINT_LOADED = "logging.checkpoint.loaded"  # was: checkpoint_loaded

    # === Configuration Events ===
    CONFIGURATION_LOADED = "config.loaded"  # was: configuration_changed
    CONFIGURATION_UPDATED = "config.updated"
    CONFIGURATION_VALIDATED = "config.validated"


class EventSeverity(Enum):
    """Event severity levels for prioritization and filtering."""

    DEBUG = "debug"
    INFO = "info"
    WARNING = "warning"
    ERROR = "error"
    CRITICAL = "critical"


class EventStatus(Enum):
    """Event processing status for lifecycle tracking."""

    CREATED = "created"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


# ========== Protocols and Interfaces ==========


@runtime_checkable
class EventDataProtocol(Protocol):
    """Protocol for event data validation and introspection."""

    event_type: EventType
    event_id: str
    timestamp: float
    source: Optional[str]

    def validate(self) -> None:
        """Validate event data integrity."""
        ...

    def to_dict(self) -> Dict[str, Any]:
        """Convert event to dictionary for serialization."""
        ...


# ========== Base Event Data with Progressive Disclosure ==========


@dataclass
class BaseEventData(ABC):
    """
    Unified base class for all event data with progressive disclosure.

    Core fields are always present, advanced features are optional.
    Follows SOLID principles with clean separation of concerns.
    """

    # === Core Required Fields ===
    event_type: EventType

    # === Core Optional Fields (Basic Usage) ===
    event_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    timestamp: float = field(default_factory=time.time)
    source: Optional[str] = None
    correlation_id: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    # === Advanced Optional Fields (Progressive Disclosure) ===
    severity: EventSeverity = EventSeverity.INFO
    status: EventStatus = EventStatus.CREATED
    tags: List[str] = field(default_factory=list)

    # === Relationship Tracking (Advanced) ===
    causation_id: Optional[str] = None  # ID of event that caused this one
    parent_event_id: Optional[str] = None

    # === Extended Context (Advanced) ===
    context: Dict[str, Any] = field(default_factory=dict)

    # === Version Control ===
    event_data_version: ClassVar[str] = "2.0.0"  # Unified version

    def __post_init__(self) -> None:
        """Post-initialization validation and setup."""
        self.validate()
        self._add_system_metadata()

    @abstractmethod
    def validate(self) -> None:
        """Validate event data integrity."""
        if not isinstance(self.event_type, EventType):
            raise ValueError(
                f"event_type must be an EventType, got {type(self.event_type)}"
            )

        if not isinstance(self.timestamp, (int, float)) or self.timestamp < 0:
            raise ValueError(
                f"timestamp must be a non-negative number, got {self.timestamp}"
            )

    def _add_system_metadata(self) -> None:
        """Add system-generated metadata."""
        if "created_at" not in self.metadata:
            self.metadata["created_at"] = datetime.fromtimestamp(
                self.timestamp
            ).isoformat()

        if "event_version" not in self.metadata:
            self.metadata["event_version"] = self.event_data_version

    def add_metadata(self, key: str, value: Any) -> None:
        """Add metadata to the event."""
        self.metadata[key] = value

    def get_metadata(self, key: str, default: Any = None) -> Any:
        """Get metadata from the event."""
        return self.metadata.get(key, default)

    def add_context(self, key: str, value: Any) -> None:
        """Add context information to the event."""
        self.context[key] = value

    def get_context(self, key: str, default: Any = None) -> Any:
        """Get context information from the event."""
        return self.context.get(key, default)

    def add_tag(self, tag: str) -> None:
        """Add a tag to the event."""
        if tag not in self.tags:
            self.tags.append(tag)

    def has_tag(self, tag: str) -> bool:
        """Check if event has a specific tag."""
        return tag in self.tags

    def to_dict(self) -> Dict[str, Any]:
        """Convert event to dictionary for serialization."""
        return {
            field.name: getattr(self, field.name)
            for field in self.__dataclass_fields__.values()
            if getattr(self, field.name) is not None
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "BaseEventData":
        """Create event from dictionary."""
        # Filter out fields that don't belong to this class
        valid_fields = {field.name for field in cls.__dataclass_fields__.values()}
        filtered_data = {k: v for k, v in data.items() if k in valid_fields}
        return cls(**filtered_data)


# Type variable for generic event handling
E = TypeVar("E", bound=BaseEventData)


# ========== Specialized Event Data Classes ==========
# These maintain the same structure but inherit the unified BaseEventData


@dataclass
class TaskEvent(BaseEventData):
    """Event data for task-related events with unified features."""

    task_id: int = -1
    task_name: Optional[str] = None
    previous_task_id: Optional[int] = None
    num_classes: Optional[int] = None

    # Framework integration
    pl_module: Optional["pl.LightningModule"] = None
    trainer: Optional["pl.Trainer"] = None
    dataloader: Optional[DataLoader] = None

    # Metrics and configuration
    task_metrics: Optional[Dict[str, float]] = None
    task_config: Optional[Dict[str, Any]] = None

    def validate(self) -> None:
        """Validate task event data."""
        super().validate()

        if not isinstance(self.task_id, int) or self.task_id < 0:
            raise ValueError(
                f"task_id must be a non-negative integer, got {self.task_id}"
            )

        if self.previous_task_id is not None:
            if not isinstance(self.previous_task_id, int) or self.previous_task_id < 0:
                raise ValueError(
                    f"previous_task_id must be a non-negative integer, got {self.previous_task_id}"
                )


@dataclass
class StrategyEvent(BaseEventData):
    """Event data for continual learning strategy interactions."""

    strategy_name: str = ""
    request_data: Optional[Dict[str, Any]] = None
    response_data: Optional[Dict[str, Any]] = None

    # Model and training context
    model: Optional[nn.Module] = None
    criterion: Optional[nn.Module] = None
    optimizer: Optional[torch.optim.Optimizer] = None

    # Batch processing data
    batch: Optional[Any] = None
    batch_idx: Optional[int] = None

    # Loss computation data
    outputs: Optional[torch.Tensor] = None
    targets: Optional[torch.Tensor] = None
    loss_components: Optional[Dict[str, torch.Tensor]] = None

    # # Task context
    # task_id: Optional[int] = None
    # current_epoch: Optional[int] = None
    # global_step: Optional[int] = None
    # Task context
    task_id: int = -1
    current_epoch: int = -1
    global_step: Optional[int] = None

    def validate(self) -> None:
        """Validate strategy event data."""
        super().validate()

        if not isinstance(self.strategy_name, str):
            raise ValueError(
                f"strategy_name must be a string, got {type(self.strategy_name)}"
            )

        if self.task_id is not None and (
            not isinstance(self.task_id, int) or self.task_id < 0
        ):
            raise ValueError(
                f"task_id must be a non-negative integer, got {self.task_id}"
            )


# ========== Fluent Interface Event Builders ==========


class EventBuilder(ABC):
    """Abstract base class for fluent event builders."""

    def __init__(self, event_type: EventType):
        self._event_type = event_type
        self._data: Dict[str, Any] = {}

    def source(self, source: str) -> Self:
        """Set the event source."""
        self._data["source"] = source
        return self

    def severity(self, severity: EventSeverity) -> Self:
        """Set the event severity."""
        self._data["severity"] = severity
        return self

    def correlation_id(self, correlation_id: str) -> Self:
        """Set the correlation ID."""
        self._data["correlation_id"] = correlation_id
        return self

    def metadata(self, metadata: Dict[str, Any]) -> Self:
        """Set event metadata."""
        self._data["metadata"] = metadata
        return self

    def add_metadata(self, key: str, value: Any) -> Self:
        """Add a single metadata entry."""
        if "metadata" not in self._data:
            self._data["metadata"] = {}
        self._data["metadata"][key] = value
        return self

    def context(self, context: Dict[str, Any]) -> Self:
        """Set event context."""
        self._data["context"] = context
        return self

    def add_context(self, key: str, value: Any) -> Self:
        """Add a single context entry."""
        if "context" not in self._data:
            self._data["context"] = {}
        self._data["context"][key] = value
        return self

    def tags(self, tags: List[str]) -> Self:
        """Set event tags."""
        self._data["tags"] = tags
        return self

    def add_tag(self, tag: str) -> Self:
        """Add a single tag."""
        if "tags" not in self._data:
            self._data["tags"] = []
        if tag not in self._data["tags"]:
            self._data["tags"].append(tag)
        return self

    def response_data(self, response_data: Dict[str, Any]) -> Self:
        """Set response data."""
        self._data["response_data"] = response_data
        return self

    @abstractmethod
    def build(self) -> BaseEventData:
        """Build the final event object."""
        pass


class TaskEventBuilder(EventBuilder):
    """Fluent builder for TaskEvent objects."""

    def task_id(self, task_id: int) -> Self:
        """Set the task ID."""
        self._data["task_id"] = task_id
        return self

    def task_name(self, task_name: str) -> Self:
        """Set the task name."""
        self._data["task_name"] = task_name
        return self

    def previous_task_id(self, previous_task_id: int) -> Self:
        """Set the previous task ID."""
        self._data["previous_task_id"] = previous_task_id
        return self

    def num_classes(self, num_classes: int) -> Self:
        """Set the number of classes."""
        self._data["num_classes"] = num_classes
        return self

    def pl_module(self, pl_module: pl.LightningModule) -> Self:
        """Set the PyTorch Lightning module."""
        self._data["pl_module"] = pl_module
        return self

    def trainer(self, trainer: pl.Trainer) -> Self:
        """Set the PyTorch Lightning trainer."""
        self._data["trainer"] = trainer
        return self

    def dataloader(self, dataloader: DataLoader) -> Self:
        """Set the dataloader."""
        self._data["dataloader"] = dataloader
        return self

    def task_metrics(self, task_metrics: Dict[str, float]) -> Self:
        """Set task metrics."""
        self._data["task_metrics"] = task_metrics
        return self

    def task_config(self, task_config: Dict[str, Any]) -> Self:
        """Set task configuration."""
        self._data["task_config"] = task_config
        return self

    def build(self) -> TaskEvent:
        """Build the TaskEvent object."""
        return TaskEvent(event_type=self._event_type, **self._data)


class StrategyEventBuilder(EventBuilder):
    """Fluent builder for StrategyEvent objects."""

    def strategy_name(self, strategy_name: str) -> Self:
        """Set the strategy name."""
        self._data["strategy_name"] = strategy_name
        return self

    def request_data(self, request_data: Dict[str, Any]) -> Self:
        """Set request data."""
        self._data["request_data"] = request_data
        return self

    def response_data(self, response_data: Dict[str, Any]) -> Self:
        """Set response data."""
        self._data["response_data"] = response_data
        return self

    def model(self, model: nn.Module) -> Self:
        """Set the model."""
        self._data["model"] = model
        return self

    def criterion(self, criterion: nn.Module) -> Self:
        """Set the criterion."""
        self._data["criterion"] = criterion
        return self

    def optimizer(self, optimizer: torch.optim.Optimizer) -> Self:
        """Set the optimizer."""
        self._data["optimizer"] = optimizer
        return self

    def batch(self, batch: Any) -> Self:
        """Set the batch data."""
        self._data["batch"] = batch
        return self

    def batch_idx(self, batch_idx: int) -> Self:
        """Set the batch index."""
        self._data["batch_idx"] = batch_idx
        return self

    def outputs(self, outputs: torch.Tensor) -> Self:
        """Set the model outputs."""
        self._data["outputs"] = outputs
        return self

    def targets(self, targets: torch.Tensor) -> Self:
        """Set the targets."""
        self._data["targets"] = targets
        return self

    def loss_components(self, loss_components: Dict[str, torch.Tensor]) -> Self:
        """Set the loss components."""
        self._data["loss_components"] = loss_components
        return self

    def task_id(self, task_id: int) -> Self:
        """Set the task ID."""
        self._data["task_id"] = task_id
        return self

    def current_epoch(self, current_epoch: int) -> Self:
        """Set the current epoch."""
        self._data["current_epoch"] = current_epoch
        return self

    def global_step(self, global_step: int) -> Self:
        """Set the global step."""
        self._data["global_step"] = global_step
        return self

    def build(self) -> StrategyEvent:
        """Build the StrategyEvent object."""
        return StrategyEvent(event_type=self._event_type, **self._data)


# ========== Enhanced Event Factory with Fluent Interface ==========


class EventFactory:
    """
    Enhanced factory for creating events with both traditional and fluent interfaces.

    Supports both simple creation and fluent builder patterns for maximum flexibility.
    """

    # ========== Traditional Static Methods (Backward Compatibility) ==========

    @staticmethod
    def create_task_event(
        event_type: EventType,
        task_id: int,
        source: Optional[str] = None,
        severity: EventSeverity = EventSeverity.INFO,
        **kwargs,
    ) -> TaskEvent:
        """Create a task event with proper validation."""
        return TaskEvent(
            event_type=event_type,
            task_id=task_id,
            source=source,
            severity=severity,
            **kwargs,
        )

    @staticmethod
    def create_strategy_event(
        event_type: EventType,
        strategy_name: str,
        source: Optional[str] = None,
        severity: EventSeverity = EventSeverity.INFO,
        **kwargs,
    ) -> StrategyEvent:
        """Create a strategy event with proper validation."""
        return StrategyEvent(
            event_type=event_type,
            strategy_name=strategy_name,
            source=source,
            severity=severity,
            **kwargs,
        )

    # ========== Fluent Interface Methods ==========

    @staticmethod
    def task(event_type: EventType) -> TaskEventBuilder:
        """Start building a task event with fluent interface."""
        return TaskEventBuilder(event_type)

    @staticmethod
    def strategy(event_type: EventType, strategy_name: str) -> StrategyEventBuilder:
        """Start building a strategy event with fluent interface."""
        builder = StrategyEventBuilder(event_type)
        return builder.strategy_name(strategy_name)

    # ========== Convenience Methods for Common Event Types ==========

    @staticmethod
    def task_started(task_id: int) -> TaskEventBuilder:
        """Create a task started event builder."""
        return EventFactory.task(EventType.TASK_STARTED).task_id(task_id)

    @staticmethod
    def task_completed(task_id: int) -> TaskEventBuilder:
        """Create a task completed event builder."""
        return EventFactory.task(EventType.TASK_COMPLETED).task_id(task_id)

    @staticmethod
    def strategy_loss_computation_completed(strategy_name: str) -> StrategyEventBuilder:
        """Create a strategy loss computation completed event builder."""
        return EventFactory.strategy(
            EventType.LOSS_COMPUTATION_COMPLETED, strategy_name
        )

    @staticmethod
    def strategy_batch_preprocessing_completed(
        strategy_name: str,
    ) -> StrategyEventBuilder:
        """Create a strategy batch preprocessing completed event builder."""
        return EventFactory.strategy(
            EventType.BATCH_PREPROCESSING_COMPLETED, strategy_name
        )

    @staticmethod
    def strategy_state_changed(strategy_name: str) -> StrategyEventBuilder:
        """Create a strategy state changed event builder."""
        return EventFactory.strategy(EventType.STRATEGY_STATE_CHANGED, strategy_name)


# ========== Type Aliases and Unions ==========

EventData = Union[
    TaskEvent,
    StrategyEvent,
    # Add other event types as they're implemented
]

# ========== Builder Type Aliases ==========

EventBuilderType = Union[
    TaskEventBuilder,
    StrategyEventBuilder,
    # Add other builders as they're implemented
]
